#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Milvus数据查看工具
功能：查看Docker Milvus中的数据
"""

from pymilvus import connections, Collection, utility


def connect_to_milvus(host="localhost", port="19530"):
    """连接到Milvus"""
    try:
        connections.connect("default", host=host, port=port)
        print(f"✅ 成功连接到Milvus服务器 {host}:{port}")
        return True
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        return False


def list_all_collections():
    """列出所有集合"""
    try:
        collections = utility.list_collections()
        print(f"\n📚 所有集合 (共{len(collections)}个):")
        
        for i, col_name in enumerate(collections, 1):
            col = Collection(col_name)
            print(f"  {i}. {col_name}: {col.num_entities} 条记录")
        
        return collections
    except Exception as e:
        print(f"❌ 获取集合列表失败: {e}")
        return []


def view_collection_data(collection_name, limit=5):
    """查看集合数据"""
    try:
        col = Collection(collection_name)
        print(f"\n🔍 查看集合: {collection_name}")
        print(f"记录总数: {col.num_entities}")
        
        # 加载集合
        col.load()
        print("集合已加载到内存")
        
        # 查询数据
        results = col.query(
            expr='id != ""', 
            limit=limit, 
            output_fields=['filename', 'text', 'timestamp']
        )
        
        print(f"\n📄 前{len(results)}条记录:")
        for i, result in enumerate(results, 1):
            print(f"\n--- 记录 {i} ---")
            print(f"文件名: {result.get('filename', 'N/A')}")
            print(f"时间: {result.get('timestamp', 'N/A')}")
            text = result.get('text', '')
            if len(text) > 200:
                print(f"内容: {text[:200]}...")
            else:
                print(f"内容: {text}")
        
        return True
        
    except Exception as e:
        print(f"❌ 查看数据失败: {e}")
        return False


def search_by_filename(collection_name, filename_pattern):
    """按文件名搜索"""
    try:
        col = Collection(collection_name)
        col.load()
        
        # 使用like查询
        expr = f'filename like "{filename_pattern}%"'
        results = col.query(
            expr=expr,
            output_fields=['filename', 'text', 'timestamp']
        )
        
        print(f"\n🔍 搜索结果 (文件名包含'{filename_pattern}'):")
        print(f"找到 {len(results)} 条记录")
        
        for i, result in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(f"文件名: {result['filename']}")
            print(f"时间: {result['timestamp']}")
            text = result['text']
            if len(text) > 150:
                print(f"内容: {text[:150]}...")
            else:
                print(f"内容: {text}")
        
        return results
        
    except Exception as e:
        print(f"❌ 搜索失败: {e}")
        return []


def main():
    """主函数"""
    print("🔍 Milvus数据查看工具")
    print("=" * 50)
    
    # 连接Milvus
    if not connect_to_milvus():
        return
    
    # 列出所有集合
    collections = list_all_collections()
    if not collections:
        return
    
    while True:
        print("\n" + "=" * 50)
        print("请选择操作:")
        print("1. 查看集合数据")
        print("2. 按文件名搜索")
        print("3. 重新列出集合")
        print("4. 退出")
        
        choice = input("请输入选择 (1-4): ").strip()
        
        if choice == "1":
            print("\n可用集合:")
            for i, col_name in enumerate(collections, 1):
                print(f"  {i}. {col_name}")
            
            try:
                col_choice = int(input("请选择集合编号: ")) - 1
                if 0 <= col_choice < len(collections):
                    limit = int(input("显示记录数 (默认5): ") or "5")
                    view_collection_data(collections[col_choice], limit)
                else:
                    print("❌ 无效选择")
            except ValueError:
                print("❌ 请输入有效数字")
        
        elif choice == "2":
            print("\n可用集合:")
            for i, col_name in enumerate(collections, 1):
                print(f"  {i}. {col_name}")
            
            try:
                col_choice = int(input("请选择集合编号: ")) - 1
                if 0 <= col_choice < len(collections):
                    pattern = input("请输入文件名模式 (如 'C1' 查找C1开头的文件): ")
                    search_by_filename(collections[col_choice], pattern)
                else:
                    print("❌ 无效选择")
            except ValueError:
                print("❌ 请输入有效数字")
        
        elif choice == "3":
            collections = list_all_collections()
        
        elif choice == "4":
            print("👋 再见!")
            break
        
        else:
            print("❌ 无效选择")


if __name__ == "__main__":
    main()
